from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("logs")

train_dataset = datasets.CIFAR10(root="datasets", train=True, transform=None, download=True)
test_dataset = datasets.CIFAR10(root="datasets", train=False, transform=None, download=True)

compose_transform = transforms.Compose([transforms.ToTensor()])
# 经过转换之后的测试数据集
test_transformed_dataset = datasets.CIFAR10(root="datasets", train=False, transform=compose_transform, download=True)
print(test_dataset[0])
print(test_transformed_dataset[0])

for i in range(100):
    img, target = test_transformed_dataset[i]
    writer.add_image("test_dataset", img, i)

print("add test_dataset complete !")

tensor_img, target = test_transformed_dataset[0]
writer.add_image("test " + test_transformed_dataset.classes[target], tensor_img, 1)
print("add single test dataset complete ")

# using dataloader
dataloader = DataLoader(dataset=test_transformed_dataset, batch_size=50, shuffle=True, num_workers=0, drop_last=False)
for epoch in range(2):
    step = 1
    for imgs, targets in dataloader:
        writer.add_images("dataset_dataloader_shuffle: {}".format(epoch), imgs, step)
        step = step + 1

dataloader = DataLoader(dataset=test_transformed_dataset, batch_size=50, shuffle=False, num_workers=0, drop_last=False)
for epoch in range(2):
    step = 1
    for imgs, targets in dataloader:
        writer.add_images("dataset_dataloader_no_shuffle: {}".format(epoch), imgs, step)
        step = step + 1

print("add images using dataloader complete !")

writer.close()